package cz.vutbr.fit.dnacompress;

import java.io.*;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;
import java.util.TreeSet;
import java.util.Map.Entry;

import org.biojava.bio.BioException;
import org.biojava.bio.alignment.IllegalAlignmentEditException;
import org.biojava.bio.program.ssbind.AnnotationFactory;
import org.biojava.bio.seq.DNATools;
import org.biojava.bio.seq.Sequence;
import org.biojava.bio.seq.db.SequenceDB;
import org.biojava.bio.seq.impl.SimpleSequence;
import org.biojava.bio.seq.io.SeqIOTools;
import org.biojava.bio.seq.io.SymbolTokenization;
import org.biojava.bio.symbol.Edit;
import org.biojava.bio.symbol.SimpleSymbolList;
import org.biojava.bio.symbol.Symbol;
import org.biojava.bio.symbol.SymbolList;
import org.biojava.utils.ParserException;

import cz.vutbr.fit.dnacompress.alignment.AbstractAlignmentMaker;
import cz.vutbr.fit.dnacompress.alignment.MegablastAlignmentMaker;
import cz.vutbr.fit.dnacompress.encoding.AbstractCoder;
import cz.vutbr.fit.dnacompress.encoding.StringCoder;
import cz.vutbr.fit.dnacompress.events.AlignDeletitionEvent;
import cz.vutbr.fit.dnacompress.events.AlignEvent;
import cz.vutbr.fit.dnacompress.events.AlignInsertionEvent;
import cz.vutbr.fit.dnacompress.events.AlignMutationEvent;

public class Compressor {
	
	/**
	 * stavy automatu pro findEvents
	 */
	private static final int	READ = 0;
    private static final int    DELETITION = 1;
    private static final int 	INSERTION = 2;
    private static final int   	MUTATION = 3;
    private static final int   	FINISH = 4;
    
    /**
     * koncovka komprimovanych souboru
     */
    private static final String compressExtension = ".cbr";
    
    /**
     * obecny koder pro kodovani i dekodovani eventu
     */
    private AbstractCoder coder = null;
   
    /**
     * tokenizer pro vytvoreni SymbolListu
     */
	private SymbolTokenization tokenParser;
	
	/**
	 * referencni sekvence
	 */
	private Sequence referenceSeq;
	
	/**
	 * nazev referencni sekvence
	 */
	private String referenceFileName;
	
	/**
	 * query identifikator
	 */
	private String queryID;
	
	/**
	 * referencni identifikator
	 */
	private String referenceID;

	/**
	 * konstruktory
	 */
	public Compressor() {}
	
	public Compressor( String referenceFileName) throws Exception {
		
		setReference( referenceFileName);
		setCoder(new StringCoder());
	}
	
	public Compressor( String referenceFileName, AbstractCoder coder) throws Exception {
		
		setReference( referenceFileName);
		setCoder( coder);
	}
	
	
	/**
	 * hlavni metoda pro komprimaci DNA
	 * 
	 * @param queryFileName
	 * @param alignmentMaker
	 * 
	 * @return boolean
	 * @throws BioException
	 */
	public boolean compress( String queryFileName, AbstractAlignmentMaker alignmentMaker) 
		throws BioException 
	{
		
		String queryBaseName = Main.basename( queryFileName);
		
		// inicialiace align-makeru
		alignmentMaker.init( this.referenceFileName, queryFileName);
		// spusteni alignmakeru
		alignmentMaker.run();
		
		// nalezeni eventu
		List<AlignEvent> events = findEvents( alignmentMaker.getReference(), alignmentMaker.getQuery());
		System.out.println( "Aligned events: " + events);
		
		// init output header
		this.setQueryID( getFastaHeader( Main.dataFolder + queryFileName));
		this.setReferenceID( referenceSeq.getName());
		
		// write data to file
		try {
			OutputStream output;
			output = new BufferedOutputStream(new FileOutputStream( Main.dataFolder + queryBaseName + compressExtension));
			// write header
			output.write( this.makeHeader( coder.getCoderID()));
			// write data
			output.write( coder.encode(events));
		    output.close();
		} catch (IOException e) {
			
			e.printStackTrace();
		}
	    
	    System.out.println("Your file has been written"); 
	    
		return true;
	}
	
	/**
	 * Hlavni metoda pro dekompresi DNA
	 * 
	 * @param String compressedFileName
	 * @return boolean
	 * @throws ParserException 
	 */
	public boolean decompress( String compressedFileName) throws ParserException {
		
		// nazev noveho souboru s DNA sekvenci
		String decompressedFileName = Main.basename( compressedFileName) + "-new.fa";
		// tabylka poctu
		String countTableFileName = Main.dataFolder + Main.basename( referenceFileName) + ".ctb";
		
		try {
			File f = new File(Main.dataFolder + compressedFileName);
			String ls = System.getProperty("line.separator");
			
			String line = null;
			int offset = 0;
			
			/**
			 * rozparsovani hlavicky
			 */
			BufferedReader br = new BufferedReader(new FileReader( f));
			// query line
			line = br.readLine();
			offset = line.getBytes().length + ls.length();
			this.setQueryID( line);
			System.out.println( "Compressed query: \t" + this.getQueryID());
			// reference line
			line = br.readLine();
			offset += line.getBytes().length + ls.length();
			this.setReferenceID( line);
			System.out.println( "Reference used: \t" + this.getReferenceID());
			// coder line
			line = br.readLine();
			offset += line.getBytes().length + ls.length();
			this.setCoder( AbstractCoder.getClass( Integer.parseInt( line), countTableFileName));
			System.out.println( "Coder used: \t" + this.getCoder());
			
			/**
			 * dekodovani dat ze souboru
			 */
			InputStream is = new FileInputStream( f);
			is.skip( offset);
			
			// get list of alignEvent from encoded data
			List<AlignEvent> events = coder.decode( is);
			
			System.out.println( "Decoded events: " + events);
			
			// copy of reference to make original
			SimpleSymbolList syms = new SimpleSymbolList(this.referenceSeq);
			Edit e = null;
			int position = 0;
			System.out.println( "Changing reference with events ...");
			for( Iterator<AlignEvent> i = events.iterator(); i.hasNext();) {
				AlignEvent event = i.next();
				// position by intervals
				position += event.getPosition();
				
				// vytvoreni zmeny referencni -> original zavysle na druhu eventu
				if( event instanceof AlignMutationEvent) {
					e = new Edit( position, event.getData().length(), DNATools.createDNA( event.getData()));
				} else if( event instanceof AlignInsertionEvent) {
					e = new Edit( position, 0, DNATools.createDNA( event.getData()));
				} else if( event instanceof AlignDeletitionEvent){
					e = new Edit( position, Integer.parseInt( event.getData()), SymbolList.EMPTY_LIST);
				} else {
					throw new IllegalAlignmentEditException( "Unexpected object in events list!");
				}
			    //apply the edit
			    syms.edit(e);
			}
			// vytvoreni sekvence ze symbollistu
			System.out.println( "Changing reference complete.");
			Sequence uncompressedSequence = new SimpleSequence( syms, "urn", this.getQueryID(), AnnotationFactory.makeAnnotation( new HashMap()));
			
			// zapis do sobuoru
			System.out.println( "Writing data to fasta format output.");
			SeqIOTools.writeFasta( new FileOutputStream( Main.dataFolder + decompressedFileName), uncompressedSequence);
			
		    
		} catch (IOException e) {
				
			e.printStackTrace();
		} catch (BioException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
		
		System.out.println( "Operation complete.");
		return true;
	}
	
	/**
	 * vytvoreni tabulky poctu ze skupiny souboru a referencni sekvence
	 * 
	 * @param path	Cesta k testovacim souborum
	 * @param referenceFileName	Nazev referencni sekvence, ktera musi byt ve slozce path
	 * @throws BioException
	 */
	public void createCountTable( String path, String referenceFileName) throws BioException {
		
		// slozka
		File dir = new File( path);
		
		// vytvoreni alignMakeru pro vytvoreni zarovnani
		MegablastAlignmentMaker am = new MegablastAlignmentMaker();
		am.setDataFolder( path);
		
		// seznam vsech eventu nalezenych ve vsech souborech
		List<AlignEvent> events = new ArrayList<AlignEvent>();
		
		// This filter only returns directories 
		File[] files = dir.listFiles();	
		// projde vsechny soubory
		for(int i=0;i<files.length;i++) {
			
			File file = files[i];
			// je to soubor a neni reference
			if( file.isFile() && !file.getName().equals( (String) referenceFileName)) {
				
				// zarovnani
				am.init(referenceFileName, file.getName(), "alignment");
				am.run();
				
				// find events and connect to event list
				events.addAll( findEvents( am.getReference(), am.getQuery()));
			}
		}
		System.out.println( events.size());
		
		/**
		 * vytvoreni mapy kde klic jsou data a data je AlignEvent object, ktery v sobe drzi pocet vyskytu
		 */
		TreeMap<String, AlignEvent> table = new TreeMap<String, AlignEvent>();
		for( Iterator<AlignEvent> i = events.iterator(); i.hasNext();) {
			
			AlignEvent event = i.next();
			String data = event.getSeparedData();
			if( !table.containsKey( data)) {
				table.put( data, event);
			}
			table.get( data).incrementCount();
		}
		
		// seradi mapu podle vyskytu
		TreeSet<Map.Entry<String, AlignEvent>> set = new TreeSet<Map.Entry<String, AlignEvent>>( 
			new Comparator<Map.Entry<String, AlignEvent>>() {
				@Override
				public int compare(Entry<String, AlignEvent> o1, Entry<String, AlignEvent> o2) {
					
					Integer count1 = o1.getValue().getCount();
					Integer count2 = o2.getValue().getCount();
					
					
					int compareTo = count2.compareTo( count1);
					if( compareTo == 0) {
						return o1.getValue().getSeparedData().compareTo( o2.getValue().getSeparedData());
					} else {
						return compareTo;
					}
				}   
			}
		);  
		set.addAll( table.entrySet());

		/**
		 * vytvoreni tabulky poctu
		 */
		try {
			String tableFileName = Main.dataFolder + Main.basename( referenceFileName) + ".ctb";
			BufferedWriter bw = new BufferedWriter( new FileWriter( tableFileName));
			
			int identifier = 1;
			for (Iterator<Map.Entry<String, AlignEvent>> i = set.iterator(); i.hasNext();) {
	            Map.Entry<String, AlignEvent> entry = i.next();
	            AlignEvent ae = entry.getValue();
	            // koduje se identifikator poctu (Mut-c) a poradi poctu
	            String row = ae.getCountID() + "|" + identifier++;
	            bw.write( row);
	            bw.newLine();
	        }
			bw.close();
		} catch (IOException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
	}
		
	
	/**
	 * Vytvoreni seznamu AlignEvent objektu z porovnani dvou retezcu pomoci automatu
	 * 
	 * @param String reference 
	 * @param String query
	 * @return List<AlignEvent>
	 * @throws BioException
	 */
	public List<AlignEvent> findEvents( String reference, String query)
		    throws BioException
	{
		List<AlignEvent> events = new ArrayList<AlignEvent>();
		
		// create SymbolTokenization
	    tokenParser = DNATools.getDNA().getTokenization("token");
        
	    // make symbolLists
        SymbolList querySymbolList = new SimpleSymbolList(tokenParser, query);
        SymbolList referenceSymbolList = new SimpleSymbolList(tokenParser, reference);
		
        // make iterators for comparition by symbol
		Iterator queryIterator = querySymbolList.iterator();
		Iterator subjectIterator = referenceSymbolList.iterator();
		
		// finite automata for finding events in aligned sequences
		int state = READ;				// state of automata
		int position = 0;				// position of event in sequence
		int lastEventPosition = 0;		// interval between two events
		AlignEvent actualEvent = null;	// active event
		Symbol subjectSymbol = null;	// actual subject symbol
		Symbol querySymbol = null;		// actual query symbol
		Symbol gap = DNATools.getDNA().getGapSymbol(); // symbol for gap
		// search changes
		for (;;) {
			
			// finite state
			if( state == FINISH) break;
			
			// automata state switcher
			switch( state) {
			
				case READ: 
					
					if( queryIterator.hasNext() && subjectIterator.hasNext()) {
						querySymbol = (Symbol) queryIterator.next();
						subjectSymbol = (Symbol) subjectIterator.next();
						
						// increment position in sequences
						position++;
						
						// symbols are equal -> nothing
						if( querySymbol == subjectSymbol) {
							state = READ;
							// actualEvent is active
							if( actualEvent != null) {
								// add actual event to list of events
								events.add( actualEvent);
								actualEvent = null;
							}
						// symbols are diferent -> event	
						} else {
							// gap in reference -> insertion
							if( subjectSymbol == gap) {
								state = INSERTION;
								// ending previous event
								if( actualEvent != null && !(actualEvent instanceof AlignInsertionEvent)) {
									events.add( actualEvent);
									actualEvent = null;
								}
							// gap in query -> deletition	
							} else if( querySymbol == gap) {
								state = DELETITION;
								// ending previous event
								if( actualEvent != null && !(actualEvent instanceof AlignDeletitionEvent)) {
									events.add( actualEvent);
									actualEvent = null;
								}
								
							// otherwise -> mutation	
							} else {
								state = MUTATION;
								// ending previous event
								if( actualEvent != null && !(actualEvent instanceof AlignMutationEvent)) {
									// add actual event to list of events
									events.add( actualEvent);
									actualEvent = null;
								}
							}
						}
					// end of iterators	
					} else {
						state = FINISH;
					}
					break;
					
				case MUTATION:
					if( actualEvent == null) {
						actualEvent = new AlignMutationEvent( position - lastEventPosition);
						lastEventPosition = position;
					}
					actualEvent.addData( tokenParser.tokenizeSymbol( querySymbol));
					state = READ;
					break;
					
				case INSERTION:
					// neni rozpracovany zadny event
					if( actualEvent == null) {
						actualEvent = new AlignInsertionEvent( position - lastEventPosition);
						lastEventPosition = position;
					}
					
					actualEvent.addData( tokenParser.tokenizeSymbol( querySymbol));
					state = READ;
					break;
					
				case DELETITION:
					if( actualEvent == null) {
						actualEvent = new AlignDeletitionEvent( position - lastEventPosition);
						lastEventPosition = position;
					}
					// decrementing posistion because of deleting in query sequence
					position--;
					
					actualEvent.addData();
					state = READ;
					break;
				default:
					throw new BioException( "Wrong state in finite automata!");
			}
		}
		
		return events;
	}
	
	/**
	 * nastavi referencni sekvenci z nazvu souboru
	 * 
	 * @param fileName
	 * @throws Exception
	 */
	public void setReference( String fileName) throws Exception {
		
		this.referenceFileName = fileName;
		this.referenceSeq = createSequenceFromFile( Main.dataFolder + fileName);
	}
	
	/**
	 * Vraci referencni sekvenci
	 * @return Sequence
	 */
	public Sequence getReference() {
		
		return this.referenceSeq;
	}

	
	/**
	 * @return the coder
	 */
	public AbstractCoder getCoder() {
		return coder;
	}

	/**
	 * @param coder the coder to set
	 */
	public void setCoder(AbstractCoder coder) {
		
		this.coder = coder;
	}
	
	/**
	 * read first line in file a return it if recognize a fasta header
	 * 
	 * @param fileName
	 * @return fasta header
	 */
	public String getFastaHeader( String fileName) {

		String header = null;
		try{
			// Get the object of DataInputStream
			DataInputStream in = new DataInputStream( new FileInputStream( fileName));
			BufferedReader br = new BufferedReader(new InputStreamReader(in));
			
			//Read first line only
			header = br.readLine();
			if( !header.startsWith(">")) {
				throw new ParserException( "No fasta header found in '" + fileName + "'!");
			}
			//Close the input stream
			in.close();
		}catch ( ParserException e) {
			System.err.println("ParseError: " + e.getMessage());
		}catch (IOException e){//Catch exception if any
			System.err.println("IOError: " + e.getMessage());
		}
		return header;
	}

	
	/**
	 * create a DNA sequence from file in fasta format
	 * 
	 * @param fileName
	 * @return sequence
	 * @throws Exception 
	 */
	public Sequence createSequenceFromFile( String fileName) throws Exception {
		
		Sequence seq = null;
		
		BufferedInputStream is = new BufferedInputStream(new FileInputStream( fileName));

		//get a SequenceDB of all sequences in the file
		SequenceDB db = SeqIOTools.readFasta(is, DNATools.getDNA());
		
		seq = db.sequenceIterator().nextSequence();
		
		return seq;
	}
	
	/**
	 * 
	 * @return String
	 */
	public String getQueryID() {
		return queryID;
	}

	/**
	 * 
	 * @param queryID
	 */
	public void setQueryID(String queryID) {
		
		if( queryID.startsWith(">")) {
			queryID = queryID.substring(1);
		}
		this.queryID = queryID;
	}

	/**
	 * 
	 * @param referenceID
	 */
	public void setReferenceID(String referenceID) {
		if( referenceID.startsWith(">")) {
			referenceID = referenceID.substring(1);
		}
		this.referenceID = referenceID;
	}

	/**
	 * 
	 * @return
	 */
	public String getReferenceID() {
		return referenceID;
	}
	
	/**
	 * encode header information to byte sequence
	 * 
	 * @return 
	 */
	public byte[] makeHeader( int coderID) {
		
		String ls = System.getProperty("line.separator");
		
		// set header
		StringBuffer header = new StringBuffer( 200);
		// append query ID
		header.append( ">" + this.getQueryID());
		header.append( ls);
		// append reference ID
		header.append( ">" + this.getReferenceID());
		header.append( ls);
		// append coder ID
		header.append( coderID);
		header.append( ls);
		
		return header.toString().getBytes();
	}

}


